package tool

import java.io.File
import implicits.Implicits._
import org.apache.commons.io.FileUtils
import tool.Pojo.CommandData
import utils.Utils

import scala.jdk.CollectionConverters._
import implicits.Implicits._

import java.nio.file.Files

/**
 * Created by yz on 10/4/2021
 */
object Tool {
  val dbName = "mummer_lib"
  val playPath = new File("./").getAbsolutePath
  val linuxPath = playPath + s"/${dbName}"
  val windowsPath = s"C:\\${dbName}"
  val genomeFileSuffix = ".genome.fa"
  val gffFileSuffix = ".gff"
  val cdsFileSuffix = ".cds.fa"

  def isWindows = {
    System.getProperty("os.name") match {
      case x if x.contains("Windows") => true
      case _ => false
    }
  }

  val pyPath = {
    val path = "C:\\workspaceForIDEA\\mummer\\pyScripts"
    val linuxPyPath = linuxPath + "/pyScripts"
    if (new File(path).exists) new File(path) else new File(linuxPyPath)
  }
  val path = {
    if (new File(windowsPath).exists()) windowsPath else linuxPath
  }
  val dataFile = new File(path, "data")
  val binPath = new File(path, "bin")
  val svBin = new File(binPath, "SV_finder_2.2.1/bin")
  val windowsTestDir = new File("C:\\mummer_database\\test")
  val linuxTestDir = new File(playPath, "workspace")
  val testDir = if (windowsTestDir.exists()) windowsTestDir else linuxTestDir

  def getGbffFiles(tmpDir: File) = {
    tmpDir.listFiles().toList.filter(_.getName.endsWith(".gbff"))
  }

  def getReadGeneWorkspaceDir(parentWorkpaceDir: File) = {
    new File(parentWorkpaceDir, s"read_gene")
  }

  def getGenePredictDir(parentWorkpaceDir: File) = {
    new File(parentWorkpaceDir, s"gene_predict")
  }

  def readGeneBanks(parentWorkspaceDir: File) = {
    val pyFile = new File(Tool.pyPath, "readGenBank.py")
    val gbffFiles = Tool.getGbffFiles(parentWorkspaceDir)
    val workspaceDir = Tool.getReadGeneWorkspaceDir(parentWorkspaceDir)
    gbffFiles.foreach(_.fileCopyToDir(workspaceDir))
    val commands = gbffFiles.map { gbffFile =>
      s"""
         |python3 ${pyFile.unixPath} -i ${gbffFile.unixPath}
         |""".stripMargin
    }
    CommandData(workspaceDir, commands)
  }

  def getGenomeFile(sampleName: String) = {
    new File(Tool.dataFile, s"${sampleName}.genome.fa")
  }

  def getGenomeFile(dir: File, prefix: String) = {
    new File(dir, Tool.getGenomeFileName(prefix))
  }

  def getCdsFile(workspaceDir: File, sampleName: String) = {
    new File(workspaceDir, s"${sampleName}.cds.fa")
  }

  def getGenomeFileName(prefix: String) = {
    s"${prefix}.genome.fa"
  }

  def createTempDirectory(prefix: String) = {
    if (Tool.isWindows) Tool.testDir else Files.createTempDirectory(prefix).toFile
  }

  def genePredicts(parentWorkspaceDir: File) = {
    val gbffFiles = Tool.getGbffFiles(parentWorkspaceDir)
    val workspaceDir = Tool.getGenePredictDir(parentWorkspaceDir)
    val readGeneDir = Tool.getReadGeneWorkspaceDir(parentWorkspaceDir)
    val commands = gbffFiles.map { gbffFile =>
      val prefix = gbffFile.namePrefix
      val geneInfoFile = new File(readGeneDir, s"${prefix}_geneInfo.txt")
      val geneInfoLines = geneInfoFile.txtLines.drop(1)
      val infoFile = new File(readGeneDir, s"${prefix}_info.txt")
      val infoLines = infoFile.txtLines.drop(1)
      val columns = infoLines.head
      val kind = columns(8)
      if (geneInfoLines.isEmpty) {
        Tool.getGenomeFile(readGeneDir, prefix).fileCopyToDir(workspaceDir)
        s"""
           |perl ${Tool.binPath.unixPath}/01.Gene-prediction/bin/gene-predict_v2.pl --SpeType B --SampleId ${prefix} --genemark --verbose -shape ${kind} --cpu 10  --run multi -outdir output ${Tool.getGenomeFile(workspaceDir, prefix).unixPath}
           |cp ${workspaceDir.unixPath}/output/final/${prefix}.cds ${workspaceDir.unixPath}/${prefix}.cds.fa
           |cp ${workspaceDir.unixPath}/output/final/${prefix}.pep ${workspaceDir.unixPath}/${prefix}.pep.fa
           |cp ${workspaceDir.unixPath}/output/final/${prefix}.gff ${workspaceDir.unixPath}/
             """.stripMargin
      } else ""
    }
    CommandData(workspaceDir, commands)
  }

  def mummer(workspaceDir: File, resultDir: File,querySampleName:String,refSampleName:String) = {
    val genePredictDir = Tool.getGenePredictDir(workspaceDir)
    val readGeneDir = Tool.getReadGeneWorkspaceDir(workspaceDir)
    List(readGeneDir, genePredictDir).foreach { dir =>
      dir.listFiles().filter { file =>
        List(Tool.genomeFileSuffix, Tool.gffFileSuffix, Tool.cdsFileSuffix).exists(y => file.getName.endsWith(y))
      }.foreach { file =>
        file.fileCopyToDir(workspaceDir)
      }
    }
    val fileStr = List(refSampleName, querySampleName).map { sampleName =>
      val file = Tool.getGenomeFile(workspaceDir, sampleName)
      file.getName
    }.mkString(" ")
    val command =
      s"""
         |perl ${Tool.svBin.unixPath}/SV_finder.pl ${fileStr} --prefix data --tfix ${refSampleName} --qfix ${querySampleName} --nocomp --outdir output --locate
         |cp output/Result/01.Synteny/Target-Query.parallel.png  ${resultDir.unixPath}/Target-Query.parallel.png
         |cp output/Result/01.Synteny/Target-Query.parallel.svg ${resultDir.unixPath}/Target-Query.parallel.svg
         |cp output/Result/01.Synteny/Target-Query.xoy.png ${resultDir.unixPath}/Target-Query.xoy.png
         |cp output/Result/01.Synteny/Target-Query.xoy.svg ${resultDir.unixPath}/Target-Query.xoy.svg
         |cp output/Result/01.Synteny/data.axt  ${resultDir.unixPath}/axt.txt
         |cp output/Result/01.Synteny/data.alg  ${resultDir.unixPath}/alg.txt
         |cp output/Result/01.Synteny/data.que.miss ${resultDir.unixPath}/query.miss.txt
         |cp output/Result/01.Synteny/data.tar.miss ${resultDir.unixPath}/target.miss.txt
         |cp output/Result/01.Synteny/data.cover ${resultDir.unixPath}/cover.txt
         |cp output/Result/01.Synteny/data.identity.list ${resultDir.unixPath}/identity_list.txt
         |cp output/Result/02.Variation/data.indel ${resultDir.unixPath}/indel.txt
         |cp output/Result/02.Variation/data.snp ${resultDir.unixPath}/snp.txt
         |cp output/Result/02.Variation/data.sv ${resultDir.unixPath}/sv.txt
         |cp output/Result/02.Variation/SV_cycle.svg ${resultDir.unixPath}/SV_cycle.svg
         |cp output/Result/02.Variation/SV_cycle.png ${resultDir.unixPath}/SV_cycle.png
         |cp output/Result/02.Variation/final.snp_indel.dis.png ${resultDir.unixPath}/final.snp_indel.dis.png
         |cp output/Result/02.Variation/final.snp_indel.dis.svg ${resultDir.unixPath}/final.snp_indel.dis.svg
         |cp output/Result/02.Variation/data.snp ${workspaceDir.unixPath}/snp.txt
         |python ${Tool.pyPath.unixPath}/dealSnp.py -q ${querySampleName}
         |perl  ${Tool.binPath.unixPath}/09.Snp/src/snp_annotation_v2.2.pl snp_out.txt ${refSampleName}.genome.fa ${refSampleName}.gff --prefix data --outdir snp_anno --code 11
         |cp  snp_anno/data.cds.axt.txt ${resultDir.unixPath}/snp.cds.axt.txt
         |cp  snp_anno/data.cds.info.xls ${resultDir.unixPath}/snp.cds.info.xls
         |cp  snp_anno/data.cds.stat.xls ${resultDir.unixPath}/snp.cds.stat.xls
         |cp  snp_anno/data.chr.stat.xls ${resultDir.unixPath}/snp.chr.stat.xls
         |cp  snp_anno/data.intergenic.info.xls ${resultDir.unixPath}/snp.intergenic.info.xls
         |cp output/Result/02.Variation/data.indel ${workspaceDir.unixPath}/indel.txt
         |python ${Tool.pyPath.unixPath}/dealIndel.py
         |perl  ${Tool.binPath.unixPath}/10.InDel/src/indel_annotation_v2.pl indel_out.txt ${refSampleName}.genome.fa ${refSampleName}.gff --prefix data --outdir indel_anno --sis --code 11
         |cp  indel_anno/data.len.svg ${resultDir.unixPath}/indel.len.svg
         |cp  indel_anno/data.cds.axt ${resultDir.unixPath}/indel.cds.axt.txt
         |cp  indel_anno/data.cdsindel.stat ${resultDir.unixPath}/indel.cdsindel.stat.xls
         |cp  indel_anno/data.cds.stat ${resultDir.unixPath}/indel.cds.stat.xls
         |cp  indel_anno/data.chr.stat ${resultDir.unixPath}/indel.chr.stat.xls
         |cp  indel_anno/data.info ${resultDir.unixPath}/indel.info.xls
         |cp  indel_anno/data.intergenic.info ${resultDir.unixPath}/indel.intergenic.info.xls
         |cp  indel_anno/data.len.dis ${resultDir.unixPath}/indel.len.dis.xls
       """.stripMargin
    CommandData(workspaceDir, List(command))
  }

  def mummerResultDeal(workspaceDir: File, resultDir: File,refSampleName:String) = {
    Tool.dealMissFile(new File(resultDir, "query.miss.txt"))
    Tool.dealMissFile(new File(resultDir, "target.miss.txt"))
    val file = new File(resultDir, "sv.txt")
    val lines = file.txtLines
    val newLines = lines.map { columns =>
      columns.init
    }
    newLines.toTxtFile(file)
    Utils.svg2png(new File(resultDir, "indel.len.svg"))
    Tool.svAnno(new File(resultDir, "sv.txt"), Tool.getCdsFile(workspaceDir, refSampleName))
  }

  def svAnno(file: File, dbFile: File) = {
    case class CDSInfo(proteinId: String, start: Int, end: Int)
    val dbLines = Utils.file2Lines(dbFile)
    val map = dbLines.filter(_.startsWith(">")).map { line =>
      val columns = line.drop(1).split("\\s+")
      val infos = columns(1).replaceAll("^locus=", "").split(":")
      (infos(0), CDSInfo(columns(0), infos(1).toInt, infos(2).toInt))
    }.groupBy(_._1).mapValues(x => x.map(_._2).sortBy(_.start))
    val lines = Utils.file2Lines(file)
    val newLines = lines.map { line =>
      val columns = line.split("\t")
      val seqName = columns(0)
      val infos = map(seqName)
      val start = columns(1).toInt
      val end = columns(2).toInt
      val proteinIds = infos.filterNot { info =>
        info.end < start || info.start > end
      }.map(_.proteinId)
      val proteinIdStr = if (proteinIds.isEmpty) "NA" else proteinIds.mkString(",")
      s"${line}\t${proteinIds.size}\t${proteinIdStr}"
    }
    Utils.lines2File(new File(file.getParent, "sv.anno.txt"), newLines)
  }

  def dealMissFile(file: File) = {
    val lines = FileUtils.readLines(file).asScala
    val newLines = lines.map { line =>
      val columns = line.split("\t")
      s"${columns(0)}\t${columns(2)}"
    }
    FileUtils.writeLines(file, newLines.asJava)
  }

}
